# -*- coding: utf-8 -*-
"""
@Time: 2021/11/22 9:24
@Author: 鹄望潇湘
@File: GaitMChannelDataSet.py
@desc: 
"""
import os
import numpy as np
from PIL import Image
from tqdm import tqdm
import torch.utils.data as data
import torch
import time
import random


class SingleGaitSequence:
    """
    步态序列类，用于标识一个ID在某种类型下的某个角度中的一个行走序列
    """
    def __init__(self, image_folder: str, id, types, angle, is_cache=True, is_preprocessed=False):
        """

        :param image_folder: 该序列的数据所在的文件夹
        :param id: 该序列所属的 ID
        :param types: 该序列的类型（如背包，正常，换衣服后）
        :param angle: 拍摄该序列的角度
        :param is_cache: 是否将数据缓存到内存中
        :param is_preprocessed: 该序列数据是原始数据(False)还是经过预处理合成后的数据(True)
        """
        self.imgs = []
        self.seq_length = 0
        self.is_cache = is_cache
        self.cache = None
        self.id = id
        self.types = types
        self.angle = angle
        self.is_preprocessed = is_preprocessed
        self.shape = None
        self.decoded = None

        if is_preprocessed is False:
            self.add_image_from_image(image_folder)
        else:
            self.add_image_from_preprocessed(image_folder)
        pass

    def add_image_from_image(self, image_folder):
        """
        使用文件夹中的图像作为序列

        :param image_folder: 存放序列图像的文件夹
        :return: None
        """
        for root, dirs, files in os.walk(image_folder):
            for file in files:
                self.imgs.append(os.path.join(root, file))
        self.seq_length = len(self.imgs)
        if self.seq_length == 0:
            print(image_folder+"为空")

    def add_image_from_preprocessed(self, image_folder):
        """
        读取预处理合并后的文件，提取序列数据

        :param image_folder: 合并文件所在的文件夹
        :return: None
        """
        self.shape = np.load(os.path.join(image_folder, "shape.npy"))
        self.imgs.append(os.path.join(image_folder, "sources.npy"))
        self.seq_length = self.shape[0]

    def __getitem__(self):
        result_image = None
        remain_channel = 30
        target_frame = random.sample(range(0, self.seq_length), remain_channel)
        target_frame.sort()

        if self.is_preprocessed:
            raw = None
            if (self.cache is None and self.is_cache is True) or self.is_cache is False:
                raw = np.load(self.imgs[0])
                if self.is_cache:
                    self.cache = raw
            else:
                raw = self.cache
            template = np.array([[128, 64, 32, 16, 8, 4, 2, 1]], dtype=np.uint8)
            raw = raw.reshape(-1, 1)
            template = template.repeat(raw.shape[0], axis=0)
            raw = raw.repeat(8, axis=1)
            result = np.bitwise_and(raw, template)/template
            result = result.reshape(-1).reshape(self.shape)
            if result.shape[0] < remain_channel:
                print("ERROR！！！！" + str(self.id))
            return (result[target_frame, :, :], self.id)

        else:
            for index in range(0, remain_channel):
                img  = self.imgs[target_frame[index]]
                temp = Image.open(img)
                temp = np.expand_dims(np.array(temp), axis=0)
                if result_image is None:
                    result_image = temp
                else:
                    result_image = np.concatenate((result_image, temp), axis=0)
            return (result_image, self.id)

    def get_all_image_data(self):
        """
        获取所有的图像帧数据，该函数只对is_preprocessed=False时有效

        :return: 所有图像序列数据
        """
        result_image = None
        if self.is_preprocessed is True:
            return result_image

        for img in self.imgs:
            temp = Image.open(img)
            temp = np.expand_dims(np.array(temp), axis=0)
            if result_image is None:
                result_image = temp
            else:
                result_image = np.concatenate((result_image, temp), axis=0)

        return result_image


class GaitMChannelDataLoader:
    """
    CASIA-B 数据集加载器，支持从官网获取到的原始数据集，也支持加载经过数据合并后的数据集
    """

    def __init__(self, image_root: str, is_preprocessed=False, use_cache=False):
        """

        :param image_root: CASIA-B 数据集存放的根目录
        :param is_preprocessed: 该数据集是否是经过预处理的（数据合并后的）还是从官方获取到的原始数据集的根目录
        :param use_cache: 是否立即缓存所有数据(将所有数据加载进内存，以提升数据加载效率），若设备内存较小，请关闭此选项
        """
        self.image_root_dir = image_root
        self.classes = []
        self.imgs = []
        self.is_preprocessed = is_preprocessed
        self.use_cache = use_cache
        self.__read_info__()

    def __read_info__(self):
        index = 0
        ids = os.listdir(self.image_root_dir)
        for id in tqdm(ids):
            if os.path.isfile(os.path.join(self.image_root_dir, id)):
                continue
            self.classes.append(id)
            typess = os.listdir(os.path.join(self.image_root_dir, id))
            for types in typess:
                angles = os.listdir(os.path.join(self.image_root_dir, id, types))
                for angle in angles:
                    temp = SingleGaitSequence(
                        os.path.join(self.image_root_dir, id, types, angle),
                        index, types, angle, is_preprocessed=self.is_preprocessed, is_cache=self.use_cache)
                    if temp.seq_length >= 30:
                        self.imgs.append(temp)
            index = index + 1
        pass

    def get_dataset_with_type_and_angle(self, type: str, angle: str) -> list:
        """
        根据类型(背包，衣服，正常等）和角度(如18度，36度等等)来筛选目标数据

        :param type: 支持 bg-01, bg-02, cl-01, cl-02, nm-01, nm-02, nm-03, nm-04, nm-05, nm-06
        :param angle: 支持 000, 018, 036, 054, 072, 090, 108, 126, 144, 162, 180
        :return: list
        """
        result = []
        for item in self.imgs:
            if item.types == type and item.angle == angle:
                result.append(item)
        return result

    def get_all_data(self) -> list:
        return self.imgs

    def get_partial_data_with_id(self, id_list):
        result = []
        for item in self.imgs:
            if item.id in id_list:
                result.append(item)
        return result


class GaitDataSet(data.Dataset):
    def __init__(self, datas: list):
        super(data.Dataset, self).__init__()
        self.datas = datas

    def __getitem__(self, index):
        image_data, id = self.datas[index].__getitem__()
        return (torch.from_numpy(image_data), id, self.datas[index].angle, self.datas[index].types)

    def __len__(self):
        return len(self.datas)


def preprocessCASIA_B(source_dir, target_dir):
    print("正在读取数据\n")
    dataset = GaitMChannelDataLoader(source_dir, is_preprocessed=False)
    print("数据读取完成，正在转换数据\n")

    template = np.array([[128, 64, 32, 16, 8, 4, 2, 1]], dtype=np.uint8)
    for item in tqdm(dataset.imgs):
        target_path = os.path.join(target_dir, str(item.id), item.types, item.angle)
        if not os.path.exists(target_path):
            os.makedirs(target_path)
        data = item.get_all_image_data()
        if data is None:
            continue
        data = data.astype(bool).astype(np.uint8)
        data_list = data.reshape(-1, 8)

        it_template = template.repeat(data_list.shape[0], axis=0)
        by = np.sum(data_list*it_template, axis=1).astype(np.uint8)

        np.save(os.path.join(target_path , "sources.npy"), by)
        np.save(os.path.join(target_path , "shape.npy"), list(data.shape))

    print("数据转换完成")


if __name__=='__main__':
    """
    # 此处程序用于将原始CASIA-B数据集中每一个人每种类型下每种角度的多张图片帧合成，以提升训练时数据的加载速度，第一个参数为原始CASIA-B数据集
    # 存放的文件夹，第二个参数是合成后的数据文件存放的位置
    # 加载合成后的数据文件时，使用 GaitMChannelDataSet("C:\\DataSet\\CASIA-B-preprocessed", is_preprocessed=True) 来加载文件，第
    # 一个参数为合成后的数据文件存放的根目录
    """
    # s1 = np.load("C:\\DataSet\\CASIA-B-preprocessed\\0\\bg-01\\000\\sources.npy")
    preprocessCASIA_B("C:\\DataSet\\CASIA-B-GaitSet", "C:\\DataSet\\CASIA-B-preprocessed")
    """
    loader = GaitMChannelDataLoader("C:\\DataSet\\CASIA-B-preprocessed", is_preprocessed=True, use_cache=False)
    imgs = loader.get_dataset_with_type_and_angle("bg-01", "000")
    for item in imgs:

        print("{}, {}, {}".format(item.id, item.types, item.angle))
        s=item.__getitem__()
    """

